import tool as tl
import os
import numpy as np
labelbar = {
    'baseline 2layer cnn': 'baseline',
    '4layer_cnn_0.7': '4-layer KD',
    '6layer_cnn_0.78': '6-layer KD',
    '8layer_cnn_0.81': '8-layer KD',
    '10layer_cnn_0.84': '10-layer KD',
    'lsr kd cnn': 'LSR KD',
    'fake noise teacher': 'FNT KD'
}
def dis_picnum(path):
    class_names = os.listdir(path)

    curvedatas = []
    label = []

    for classname in class_names:
        # if( not 'cnn kd cnn' in classname) and not('baseline' in classname):
        #     continue
        # if classname == 'cnn kd cnn' or classname == 'lsr kd cnn':
        #     continue
        if (not 'fake' in classname) and (not '6' in classname):
            continue
        label.append(labelbar[classname])
        print('{} is reading...'.format(classname))
        classpath = path + classname + '/'
        resultpath = classpath + 'result/'
        result_lists = os.listdir(resultpath)
        resultmatrix = []

        for i in range(len(result_lists)):
            res = tl.to_float_dim1(tl.read_csv(resultpath + result_lists[i])[0])
            resultmatrix.append(res)

        ave_res = tl.get_mean_arr(np.array(resultmatrix))


        curvedatas.append(ave_res)
        print(ave_res[0])

        print('{} read finish'.format(classname))
    return curvedatas,label

#
# path = '../weights_params/IE demo val loss 160 epoch/'
# tl.pic_make_matrix(curvedatas, label, '', 0)
